# coding=utf-8
import math
import numpy as np
import tensorflow as tf
slim = tf.contrib.slim
import a11_net_params as netparam
"""
Compute the default anchor boxes, given an image shape.
计算出每一feature maps层上所有的prior box，格式:y,x,h,w. 这部分
的目的是为了对样本及其bbox进行处理
"""
def anchors(dtype=np.float32):
	return ssd_anchors_all_layers(
		netparam.default_params.img_shape,
		netparam.default_params.feat_shapes,
		netparam.default_params.anchor_sizes,
		netparam.default_params.anchor_ratios,
		netparam.default_params.anchor_steps,
		netparam.default_params.anchor_offset,
		dtype)

"""
Compute anchor boxes for all feature layers.
"""
def ssd_anchors_all_layers(img_shape,
                           layers_shape,     # size of the feature maps
                           anchor_sizes,     # size of the prior box
                           anchor_ratios,
                           anchor_steps,
                           offset=0.5,
                           dtype=np.float32):
    layers_anchors = []
    for i, s in enumerate(layers_shape):
        anchor_bboxes = ssd_anchor_one_layer(img_shape, 
                                             s,
                                             i,
                                             anchor_sizes[i],
                                             anchor_ratios[i],
                                             anchor_steps[i],
                                             offset=offset, 
                                             dtype=dtype)
        layers_anchors.append(anchor_bboxes)
    return layers_anchors

"""
Computer SSD default anchor boxes for one feature layer.

Determine the relative position grid of the centers, and the relative
width and height.

Arguments:
  feat_shape: Feature shape, used for computing relative position grids;
  size: Absolute reference sizes;
  ratios: Ratios to use on these features;
  img_shape: Image shape, used for computing height, width relatively to the
    former;
  offset: Grid offset.

Return:
  y, x, h, w: Relative x and y grids, and height and width.

理解：
例如第一个feature maps，其size为(38, 38), 则每一个anchor box的中心点坐标为
(y, x), y取值0-37, x取值0-37，不过函数进行了归一化处理; h和w长度都为4,记录
每个为(y,x)中心的anchor box的size，(h[0], w[0]), (h[1], w[1]) ...
这样就确定了本层的38×38×4个anchor boxes的位置
"""
def ssd_anchor_one_layer(img_shape,
                         feat_shape,
                         fsI,
                         sizes, # Smin and Smax
                         ratios,
                         step,
                         offset=0.5,
                         dtype=np.float32):
    # Compute the position grid: simple way.
    # y, x = np.mgrid[0:feat_shape[0], 0:feat_shape[1]]
    # y = (y.astype(dtype) + offset) / feat_shape[0]
    # x = (x.astype(dtype) + offset) / feat_shape[1]
    # Weird SSD-Caffe computation using steps values...
    y, x = np.mgrid[0:feat_shape[0], 0:feat_shape[1]]
    y = (y.astype(dtype) + offset) * step / img_shape[0] 
    x = (x.astype(dtype) + offset) * step / img_shape[1] 
    # Expand dims to support easy broadcasting.
    y = np.expand_dims(y, axis=-1)
    x = np.expand_dims(x, axis=-1)
    # Compute relative height and width.
    # Tries to follow the original implementation of SSD for the order.
    num_anchors = len(sizes) + len(ratios)
    h = np.zeros((num_anchors, ), dtype=dtype)
    w = np.zeros((num_anchors, ), dtype=dtype)
    # Add first anchor boxes with ratio=1.
    h[0] = sizes[0] / img_shape[0]
    w[0] = sizes[0] / img_shape[1]
    di = 1
    if len(sizes) > 1:
        h[1] = math.sqrt(sizes[0] * sizes[1]) / img_shape[0]
        w[1] = math.sqrt(sizes[0] * sizes[1]) / img_shape[1]
        di += 1
    for i, r in enumerate(ratios):
        h[i+di] = sizes[0] / img_shape[0] / math.sqrt(r)
        w[i+di] = sizes[0] / img_shape[1] * math.sqrt(r)

    return y, x, h, w

def demo_ssd_anchor_one_layer():
	y, x, h, w = ssd_anchor_one_layer (
	img_shape = [300, 300],
	feat_shape = (38, 38),
	fsI = 0, 
	sizes = (21, 45), # Smin, Smax
	ratios = [2., .5], 
	step = 8,
	offset = 0.5,
	dtype = np.float32)
	print y.shape
	print x.shape
	print h.shape
	print w.shape

if __name__ == '__main__':
	anchors();
	demo_ssd_anchor_one_layer()



